#coding=utf8
# Copyright (c) 2016 Tinydot. inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.


def simple_cluster(data_list, min_dis, weights=None):
    # data_list: [s, c]
    # min_dis: [c]
    import torch
    assert len(data_list.shape) == 2, "data list shape need be: [s, c], but:%s find"%str(data_list.shape)
    dsize, ddim = data_list.shape
    device = data_list.device
    if sum(min_dis) <= 0:
        return torch.arange(dsize).to(device).int(), 0
    if dsize < 1:
        return data_list.int(), 0
    # calc diff
    diff = (data_list.repeat(dsize, 1) - data_list.repeat(1, dsize).view(-1, ddim)).abs()
    if weights is not None:
        assert weights.shape[0] == dsize
        flat_weights = (weights.repeat(dsize, 1) + weights.repeat(1, dsize).view(-1, 1)).abs()/2.
        diff = diff/flat_weights
    diff = (diff < min_dis).min(-1)[0].view(dsize, dsize)
    bg_msk = ~torch.logical_and((torch.eye(dsize)==0).to(device).tril(), diff)
    labels = torch.arange(dsize*dsize).view(dsize, dsize).int().to(device)
    labels[bg_msk] = 0
    labels_old = labels.clone()
    # cluster merge loop
    loop_count = 0
    while True:
        labels[:] = labels.max(-2)[0]
        labels[bg_msk] = 0
        labels = labels.permute(1, 0)
        labels[:] = labels.max(-2)[0]
        labels = labels.permute(1, 0)
        labels[bg_msk] = 0
        if torch.equal(labels_old, labels):
            break
        labels_old[:] = labels
        loop_count += 1
    # select label
    clst_labels = torch.stack([labels.max(-1)[0], labels.max(-2)[0]]).max(-2)[0]
    ret_labels = (torch.arange(dsize).to(device) + torch.max(clst_labels) + 1).int()
    ulabel_msk = clst_labels > 0
    ret_labels[ulabel_msk] = clst_labels[ulabel_msk]
    return ret_labels, loop_count
